
#include "RtBase.h"
#include "RtHttpNTLMAuth.h"
#include "RtHttpURL.h"
#include "RtChannelHttpClient.h"

#include <rpc.h>
#include <lm.h>

#define SEC_SUCCESS(Status) ((Status) >= 0)
#define STATUS_SUCCESS 0

// Target name for the security package
#define TOKEN_SOURCE_NAME       "InetSvcs"

LPCSTR CRtHttpNTLMAuth::s_pszNTLMAuth = "NTLM";
HINSTANCE CRtHttpNTLMAuth::s_hSecLib;
CRtHttpNTLMAuth::SEC_FUNC CRtHttpNTLMAuth::s_sfProcs;

CRtHttpNTLMAuth::CRtHttpNTLMAuth()
	: CRtHttpAuthSteps(2)
{
	m_asAuth._fNewConversation = TRUE;
	m_asAuth._fHaveCredHandle = FALSE;
	m_asAuth._fHaveCtxtHandle = FALSE;
	m_asAuth._fUUEncodeData = TRUE;
}

CRtHttpNTLMAuth::~CRtHttpNTLMAuth()
{
	if (m_asAuth._fHaveCredHandle) {
		s_sfProcs.pFreeCredentialsHandle(&m_asAuth._hcred);
		m_asAuth._fHaveCredHandle = FALSE;
		m_asAuth._fHaveCtxtHandle = FALSE;
	}
}

CRtString CRtHttpNTLMAuth::
do_GenerateCredentials(IRtChannelHttpClient *aChannel, LPCSTR aChallenge, 
					LPCSTR aUsername, LPCSTR aPassword)
{
	CRtString strRetval;
	RT_ASSERTE_RETURN(aChannel && aChallenge && aUsername, strRetval);

	RT_INFO_TRACE_THIS("CRtHttpNTLMAuth::do_GenerateCredentials,"
		" aUsername=" << aUsername << 
// 		" aPassword=" << aPassword << 
		" aChallenge=" << aChallenge);
	if (strncasecmp(aChallenge, s_pszNTLMAuth, strlen(s_pszNTLMAuth)))
		return strRetval;

	if (!s_hSecLib && !InitAuthorizationHeader())
		return strRetval;

	static char s_szDefaultDomain[MAX_COMPUTERNAME_LENGTH+1];
	SECURITY_STATUS ss;
	TimeStamp Lifetime;
	ULONG ContextAttributes;
	if (!*s_szDefaultDomain) {
		DWORD cbN = sizeof(s_szDefaultDomain);
		::GetComputerName(s_szDefaultDomain, &cbN );
	}
	
	CRtString sbBuffIn;
	LPCSTR pBuf = aChallenge + strlen(s_pszNTLMAuth);
	while (*pBuf == ' ')
		pBuf++;
	if (*pBuf)
		RtBase64Decode(pBuf, sbBuffIn);

	if (m_asAuth._fNewConversation) {
        SEC_WINNT_AUTH_IDENTITY   AuthIdentity;
		memset( &AuthIdentity, 0, sizeof( AuthIdentity ));
		
		PVOID pAuthData = NULL;
		if (*aUsername) {
			LPCSTR pDomain = strchr(aUsername, '\\');
			if (pDomain) {
				AuthIdentity.Domain       = (unsigned char *) aUsername;
				AuthIdentity.DomainLength = pDomain - aUsername;
				AuthIdentity.User       = (unsigned char *) pDomain+1;
				AuthIdentity.UserLength = strlen(pDomain+1);
			} else {
				AuthIdentity.Domain       = (unsigned char *) s_szDefaultDomain;
				AuthIdentity.DomainLength = strlen(s_szDefaultDomain);
				AuthIdentity.User       = (unsigned char *) aUsername;
				AuthIdentity.UserLength = strlen(aUsername);
			}
			if (aPassword) {
				AuthIdentity.Password       = (unsigned char *) aPassword;
				AuthIdentity.PasswordLength = strlen(aPassword);
			}
			AuthIdentity.Flags = SEC_WINNT_AUTH_IDENTITY_ANSI;
			pAuthData = &AuthIdentity;
		}

		ss = s_sfProcs.pAcquireCredentialsHandle( 
					NULL,								// New principal
					const_cast<char*>(s_pszNTLMAuth),	// Package name
					SECPKG_CRED_OUTBOUND,
					NULL,								// Logon ID
					pAuthData,							// Auth Data
					NULL,								// Get key func
					NULL,								// Get key arg
					&m_asAuth._hcred,
					&Lifetime );
        if ( ss != STATUS_SUCCESS ) {
            RT_ERROR_TRACE_THIS("CRtHttpNTLMAuth::do_GenerateCredentials, AcquireCredentialsHandle failed! ss=" << ss);
            return strRetval;
        }
		m_asAuth._fHaveCredHandle = TRUE;
		m_asAuth._cbMaxToken = s_sfProcs.dwMaxToken;
	} else if (sbBuffIn.empty())
		return strRetval;
	
	SecBufferDesc	OutBuffDesc;
    SecBufferDesc	InBuffDesc;
	SecBuffer		OutSecBuff;
	SecBuffer		InSecBuff;
	CRtString		sbBufOut;
	sbBufOut.resize(m_asAuth._cbMaxToken);

	OutBuffDesc.ulVersion = 0;
	OutBuffDesc.cBuffers  = 1;
	OutBuffDesc.pBuffers  = &OutSecBuff;
	OutSecBuff.cbBuffer   = m_asAuth._cbMaxToken;
	OutSecBuff.BufferType = SECBUFFER_TOKEN;
	OutSecBuff.pvBuffer   = const_cast<char*>(sbBufOut.c_str());
	if (!sbBuffIn.empty()) {
		InBuffDesc.ulVersion = 0;
		InBuffDesc.cBuffers  = 1;
		InBuffDesc.pBuffers  = &InSecBuff;
		InSecBuff.cbBuffer   = sbBuffIn.length();
		InSecBuff.BufferType = SECBUFFER_TOKEN;
		InSecBuff.pvBuffer   = const_cast<char*>(sbBuffIn.c_str());
	}

	ss = s_sfProcs.pInitializeSecurityContext( 
		&m_asAuth._hcred,
		m_asAuth._fNewConversation ? NULL : &m_asAuth._hctxt,
		TOKEN_SOURCE_NAME,
		0,
		0,
		SECURITY_NATIVE_DREP,
		m_asAuth._fNewConversation ? NULL : &InBuffDesc,
		0,
		&m_asAuth._hctxt,
		&OutBuffDesc,
		&ContextAttributes,
		&Lifetime );
    
	if (!SEC_SUCCESS(ss)) {
		if (ss == SEC_E_LOGON_DENIED)
			ss = ERROR_LOGON_FAILURE;
		RT_ERROR_TRACE_THIS("CRtHttpNTLMAuth::do_GenerateCredentials, InitializeSecurityContext failed! ss=" << ss);
		return strRetval;
    }
	if (ss == SEC_I_COMPLETE_NEEDED || ss == SEC_I_COMPLETE_AND_CONTINUE) {
		if (s_sfProcs.pCompleteAuthToken != NULL) {
			ss = s_sfProcs.pCompleteAuthToken(&m_asAuth._hctxt, &OutBuffDesc);
			if (!SEC_SUCCESS(ss)) {
				RT_ERROR_TRACE_THIS("CRtHttpNTLMAuth::do_GenerateCredentials, CompleteAuthToken failed! ss=" << ss);
				return strRetval;
			}
        } else
			return strRetval;
    }
	
	CRtString strEncode(s_pszNTLMAuth);
	strEncode += ' ';
	if (OutSecBuff.cbBuffer) {
		if (m_asAuth._fUUEncodeData) {
			const BYTE* szBufTmp = static_cast<const BYTE*>(OutSecBuff.pvBuffer);
			CRtString sbBufEncode;
			RtBase64Encode(szBufTmp, OutSecBuff.cbBuffer, sbBufEncode);
			strEncode.append(sbBufEncode);
		} else {
			LPCSTR szBufTmp = static_cast<LPCSTR>(OutSecBuff.pvBuffer);
			strEncode.append(szBufTmp, OutSecBuff.cbBuffer);
        }
    }

	if (!strEncode.empty()) {
		strRetval = strEncode;
		m_asAuth._fNewConversation = FALSE;
	}
	return strRetval;
}

BOOL CRtHttpNTLMAuth::AreCredentialsReusable()
{
	return FALSE;
}

BOOL CRtHttpNTLMAuth::AllowFirstTryWithoutUserName()
{
	return TRUE;
}

LPCSTR CRtHttpNTLMAuth::GetAuthenticateScheme()
{
	return s_pszNTLMAuth;
}

BOOL CRtHttpNTLMAuth::InitAuthorizationHeader()
{
	if (s_hSecLib)
		return TRUE;

	SecPkgInfo *pspkg;
	SECURITY_STATUS ss;
    OSVERSIONINFO VerInfo;
    LPCSTR lpszDLL;
    //
    //  Find out which security DLL to use, depending on 
    //  whether we are on NT or Win95
    //
    VerInfo.dwOSVersionInfoSize = sizeof (OSVERSIONINFO);
    if (!GetVersionEx (&VerInfo)) {
        return FALSE;
    }

	if (VerInfo.dwPlatformId == VER_PLATFORM_WIN32_NT) {
		lpszDLL = "security.dll";
	} else if (VerInfo.dwPlatformId == VER_PLATFORM_WIN32_WINDOWS) {
		lpszDLL = "secur32.dll";
	} else {
		return FALSE;
	}

    s_hSecLib = LoadLibrary(lpszDLL);
    if (s_hSecLib == NULL) {
		RT_ERROR_TRACE_THIS("CRtHttpNTLMAuth::InitAuthorizationHeader, LoadLibrary failed! lpszDLL=" << lpszDLL);
		return FALSE;
    }

    s_sfProcs.pFreeCredentialsHandle 
            = (FREE_CREDENTIALS_HANDLE_FN) GetProcAddress( 
                    s_hSecLib, 
                    "FreeCredentialsHandle" );

    s_sfProcs.pQuerySecurityPackageInfo 
            = (QUERY_SECURITY_PACKAGE_INFO_FN) GetProcAddress( 
                    s_hSecLib, 
                    "QuerySecurityPackageInfoA" );

    s_sfProcs.pAcquireCredentialsHandle 
            = (ACQUIRE_CREDENTIALS_HANDLE_FN) GetProcAddress(
                    s_hSecLib, 
                    "AcquireCredentialsHandleA" );

    s_sfProcs.pFreeContextBuffer 
            = (FREE_CONTEXT_BUFFER_FN) GetProcAddress( 
                    s_hSecLib, 
                    "FreeContextBuffer" );

    s_sfProcs.pInitializeSecurityContext
            = (INITIALIZE_SECURITY_CONTEXT_FN) GetProcAddress( 
                    s_hSecLib, 
                    "InitializeSecurityContextA" );

    s_sfProcs.pCompleteAuthToken 
            = (COMPLETE_AUTH_TOKEN_FN) GetProcAddress( 
                    s_hSecLib, 
                    "CompleteAuthToken" );

    s_sfProcs.pEnumerateSecurityPackages 
            = (ENUMERATE_SECURITY_PACKAGES_FN) GetProcAddress( 
                    s_hSecLib, 
                    "EnumerateSecurityPackagesA" );

    if ( s_sfProcs.pFreeCredentialsHandle == NULL
            || s_sfProcs.pQuerySecurityPackageInfo == NULL
            || s_sfProcs.pAcquireCredentialsHandle == NULL
            || s_sfProcs.pFreeContextBuffer == NULL
            || s_sfProcs.pInitializeSecurityContext == NULL 
            || s_sfProcs.pEnumerateSecurityPackages == NULL )
    {
		RT_ERROR_TRACE_THIS("CRtHttpNTLMAuth::InitAuthorizationHeader, LoadLibrary failed! lpszDLL=" << lpszDLL);
		goto fail;
    }

	ss = s_sfProcs.pQuerySecurityPackageInfo( 
                        const_cast<char*>(s_pszNTLMAuth),
                        &pspkg );
	if ( ss != STATUS_SUCCESS ) {
		RT_ERROR_TRACE_THIS("CRtHttpNTLMAuth::InitAuthorizationHeader, QuerySecurityPackageInfo failed! ss=" << ss);
		goto fail;
	}
	s_sfProcs.dwMaxToken = pspkg->cbMaxToken;
	s_sfProcs.pFreeContextBuffer(pspkg);
	return TRUE;

fail:
	FreeLibrary( s_hSecLib );
	s_hSecLib = NULL;
	return FALSE;
}
